from pyvis.network import Network
import sqlparse
from sqlparse.sql import Identifier, Comparison, Parenthesis

def sql_to_ast(sql):
    """将SQL解析为AST节点列表"""
    parsed = sqlparse.parse(sql)[0]
    nodes = []
    node_id = 0
    id_map = {}

    def _traverse(token, parent_id=None, depth=0):
        nonlocal node_id
        current_id = node_id
        node_type = type(token).__name__
        
        # 节点特征（可根据需要扩展）
        label = f"{node_type}\n{token.value[:20]}" if token.value else node_type
        nodes.append({
            "id": current_id,
            "label": label,
            "title": str(token),
            "level": depth,
            "group": node_type
        })
        
        if parent_id is not None:
            edges.append((parent_id, current_id))
        
        node_id += 1
        if hasattr(token, 'tokens'):
            for child in token.tokens:
                _traverse(child, current_id, depth+1)

    edges = []
    _traverse(parsed)
    return nodes, edges

def visualize_diff(sql1, sql2, output_file='diff.html'):
    """可视化两个SQL的AST差异"""
    net = Network(height="800px", width="100%", directed=True)
    
    # 解析AST并添加节点
    nodes1, edges1 = sql_to_ast(sql1)
    nodes2, edges2 = sql_to_ast(sql2)
    
    # 添加第一个AST的节点（红色）
    for node in nodes1:
        net.add_node(
            n_id=node["id"],
            label=node["label"],
            title=node["title"],
            color="#FF6B6B",
            level=node["level"],
            group=node["group"]
        )
    
    # 添加第二个AST的节点（蓝色）
    offset = max(n["id"] for n in nodes1) + 1 if nodes1 else 0
    for node in nodes2:
        net.add_node(
            n_id=node["id"] + offset,
            label=node["label"],
            title=node["title"],
            color="#4D96FF",
            level=node["level"],
            group=node["group"]
        )
    
    # 添加边
    for src, dst in edges1:
        net.add_edge(src, dst, color="#FF6B6B")
    for src, dst in edges2:
        net.add_edge(src + offset, dst + offset, color="#4D96FF")
    
    # 生成交互式HTML
    net.show_buttons(filter_=['physics'])
    net.save_graph(output_file)
    print(f"可视化结果已保存到 {output_file}")

# 使用示例
sql_original = "SELECT id, name FROM users WHERE age > 18"
sql_modified = "SELECT user_id, fullname FROM customers WHERE years > 18"
visualize_diff(sql_original, sql_modified)